package com.eonmind.ai.core.chat.advisors.token;

import com.eonmind.ai.mapper.AiConversationMapper;
import com.eonmind.common.contract.exception.BusinessException;
import com.eonmind.common.provider.dto.VipDto;
import jakarta.annotation.Resource;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.advisor.api.AdvisedRequest;
import org.springframework.ai.chat.client.advisor.api.AdvisedResponse;
import org.springframework.ai.chat.client.advisor.api.StreamAroundAdvisor;
import org.springframework.ai.chat.client.advisor.api.StreamAroundAdvisorChain;
import org.springframework.ai.chat.metadata.Usage;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;

@AllArgsConstructor
@Slf4j
@Component
public class TokenCollect implements StreamAroundAdvisor {
    @Resource
    private StringRedisTemplate stringRedisTemplate;

    @Resource(name = "aiDefaultTaskExecutor")
    private ThreadPoolTaskExecutor aiDefaultTaskExecutor;

    @Resource
    private AiConversationMapper aiConversationMapper;

    @Override
    public Flux<AdvisedResponse> aroundStream(AdvisedRequest advisedRequest, StreamAroundAdvisorChain chain) {
        log.info("TokenCollect-advisedRequest：{}", advisedRequest);
        Flux<AdvisedResponse> advisedResponseFlux = chain.nextAroundStream(advisedRequest);

        return advisedResponseFlux.map(t -> {
            Usage usage = t.response().getMetadata().getUsage();
            String id = t.adviseContext().get(CHAT_MEMORY_CONVERSATION_ID_KEY).toString();
            // 异步执行token收集
            executeTokenCollect(id, usage);
            return t;
        }).doFinally(signalType -> {
            log.info("TokenCollect-signalType：{}", signalType);
        });
    }//2409c4b2-a4ae-9e04-8ca9-646486814f53

    @Override
    public String getName() {
        return "TokenCollect";
    }

    @Override
    public int getOrder() {
        return 2;
    }

    public void executeTokenCollect(String id, Usage usage) {
        aiDefaultTaskExecutor.execute(() -> {
            // 更新数据
            aiConversationMapper.incrementTotalToken(Long.parseLong(id), usage.getTotalTokens());

        });
    }

}
